      1  import java.util.ArrayList;

      2  import java.util.List;

      3

      4  /**

      5     This class implements a red-black tree whose

      6     nodes hold objects that implement the Comparable

      7     interface.

      8  */

      9  public class RedBlackTree

     10  {

     11     Node root; // Package access, for testing

     12

     13     static final int BLACK = 1;

     14     static final int RED = 0;

     15     private static final int NEGATIVE_RED = -1;

     16     private static final int DOUBLE_BLACK = 2;

     17

     18     /**

     19        Constructs an empty tree.

     20     */

     21     public RedBlackTree()

     22     {

     23        root = null;

     24     }

     25

     26     /**

     27        Inserts a new node into the tree.

     28        @param obj the object to insert

     29     */

     30     public void add(Comparable obj)

     31     {

     32        Node newNode = new Node();

     33        newNode.data = obj;

     34        newNode.left = null;

     35        newNode.right = null;

     36        if (root == null) { root = newNode; }

     37        else { root.addNode(newNode); }

     38        fixAfterAdd(newNode);

     39     }

     40

     41     /**

     42        Tries to find an object in the tree.

     43        @param obj the object to find

     44        @return true if the object is contained in the tree

     45     */

     46     public boolean find(Comparable obj)

     47     {

     48        Node current = root;

     49        while (current != null)

     50        {

     51           int d = current.data.compareTo(obj);

     52           if (d == 0) { return true; }

     53           else if (d > 0) { current = current.left; }

     54           else { current = current.right; }

     55        }

     56        return false;

     57     }

     58

     59     /**

     60        Tries to remove an object from the tree. Does nothing

     61        if the object is not contained in the tree.

     62        @param obj the object to remove

     63     */

     64     public void remove(Comparable obj)

     65     {

     66        // Find node to be removed

     67

     68        Node toBeRemoved = root;

     69        boolean found = false;

     70        while (!found && toBeRemoved != null)

     71        {

     72           int d = toBeRemoved.data.compareTo(obj);

     73           if (d == 0) { found = true; }

     74           else

     75           {

     76              if (d > 0) { toBeRemoved = toBeRemoved.left; }

     77              else { toBeRemoved = toBeRemoved.right; }

     78           }

     79        }

     80

     81        if (!found) { return; }

     82

     83        // toBeRemoved contains obj

     84

     85        // If one of the children is empty, use the other

     86

     87        if (toBeRemoved.left == null || toBeRemoved.right == null)

     88        {

     89           Node newChild;

     90           if (toBeRemoved.left == null) { newChild = toBeRemoved.right; }

     91           else { newChild = toBeRemoved.left; }

     92

     93           fixBeforeRemove(toBeRemoved);

     94           replaceWith(toBeRemoved, newChild);

     95           return;

     96        }

     97

     98        // Neither subtree is empty

     99

    100        // Find smallest element of the right subtree

    101

    102        Node smallest = toBeRemoved.right;

    103        while (smallest.left != null)

    104        {

    105           smallest = smallest.left;

    106        }

    107

    108        // smallest contains smallest child in right subtree

    109

    110        // Move contents, unlink child

    111

    112        toBeRemoved.data = smallest.data;

    113        fixBeforeRemove(smallest);

    114        replaceWith(smallest, smallest.right);

    115     }

    116

    117     /**

    118        Yields the contents of the tree in sorted order

    119        @return all data items traversed in inorder, with a space after each item

    120     */

    121     public String toString()

    122     {

    123        return toString(root);

    124     }

    125

    126

    127     /**

    128        Yields the contents of the subtree with the given root in sorted order

    129        @param parent the root of the subtree

    130        @return all data items traversed in inorder, with a space after each item

    131     */

    132     private static String toString(Node parent)

    133     {

    134        if (parent == null) { return ""; }

    135        return toString(parent.left) + parent.data + " " + toString(parent.right);

    136     }

    137

    138     /**

    139        A node of a red-black tree stores a data item and references

    140        of the child nodes to the left and to the right. The color

    141        is the "cost" of passing the node; 1 for black or 0 for red.

    142        Temporarily, it may be set to 2 or -1.

    143     */

    144     static class Node

    145     {

    146        public Comparable data;

    147        public Node left;

    148        public Node right;

    149        public Node parent;

    150        public int color;

    151

    152        /**

    153           Constructs a red node with no data.

    154        */

    155        public Node() {}

    156

    157        /**

    158           Sets the left child and updates its parent reference.

    159           @param child the new left child

    160        */

    161        public void setLeftChild(Node child)

    162        {

    163           left = child;

    164           if (child != null) { child.parent = this; }

    165        }

    166

    167        /**

    168           Sets the right child and updates its parent reference.

    169           @param child the new right child

    170        */

    171        public void setRightChild(Node child)

    172        {

    173           right = child;

    174           if (child != null) { child.parent = this; }

    175        }

    176

    177        /**

    178           Inserts a new node as a descendant of this node.

    179           @param newNode the node to insert

    180        */

    181        public void addNode(Node newNode)

    182        {

    183           int comp = newNode.data.compareTo(data);

    184           if (comp < 0)

    185           {

    186              if (left == null)

    187              {

    188                 left = newNode;

    189                 left.parent = this;

    190              }

    191              else { left.addNode(newNode); }

    192           }

    193           else if (comp > 0)

    194           {

    195              if (right == null)

    196              {

    197                 right = newNode;

    198                 right.parent = this;

    199              }

    200              else { right.addNode(newNode); }

    201           }

    202        }

    203     }

    204

    205     /**

    206        Updates the parent's and replacement node's links when this node is replaced.

    207        Also updates the root reference if it is replaced.

    208        @param toBeReplaced the node that is to be replaced

    209        @param replacement the node that replaces that node

    210     */

    211     private void replaceWith(Node toBeReplaced, Node replacement)

    212     {

    213        if (toBeReplaced.parent == null)

    214        {

    215           replacement.parent = null;

    216           root = replacement;

    217        }

    218        else if (toBeReplaced == toBeReplaced.parent.left)

    219        {

    220           toBeReplaced.parent.setLeftChild(replacement);

    221        }

    222        else

    223        {

    224           toBeReplaced.parent.setRightChild(replacement);

    225        }

    226     }

    227

    228     /**

    229        Restores the tree to a red-black tree after a node has been added.

    230        @param newNode the node that has been added

    231     */

    232     private void fixAfterAdd(Node newNode)

    233     {

    234        if (newNode.parent == null)

    235        {

    236           newNode.color = BLACK;

    237        }

    238        else

    239        {

    240           newNode.color = RED;

    241           if (newNode.parent.color == RED) { fixDoubleRed(newNode); }

    242        }

    243     }

    244

    245     /**

    246       Fixes the tree so that it is a red-black tree after a node has been removed.

    247       @param toBeRemoved the node that is to be removed

    248     */

    249     private void fixBeforeRemove(Node toBeRemoved)

    250     {

    251        if (toBeRemoved.color == RED) { return; }

    252

    253        if (toBeRemoved.left != null || toBeRemoved.right != null) // It is not a leaf

    254        {

    255           // Color the child black

    256           if (toBeRemoved.left == null) { toBeRemoved.right.color = BLACK; }

    257           else { toBeRemoved.left.color = BLACK; }

    258        }

    259        else { bubbleUp(toBeRemoved.parent); }

    260     }

    261

    262     /**

    263        Move a charge from two children of a parent

    264        @param parent a node with two children, or null (in which case nothing is done)

    265     */

    266     private void bubbleUp(Node parent)

    267     {

    268        if (parent == null) { return; }

    269        parent.color++;

    270        parent.left.color--;

    271        parent.right.color--;

    272

    273        if (bubbleUpFix(parent.left)) { return; }

    274        if (bubbleUpFix(parent.right)) { return; }

    275

    276        if (parent.color == DOUBLE_BLACK)

    277        {

    278           if (parent.parent == null) { parent.color = BLACK; }

    279           else { bubbleUp(parent.parent); }

    280        }

    281     }

    282

    283     /**

    284        Fixes a negative-red or double-red violation introduced by bubbling up.

    285        @param child the child to check for negative-red or double-red violations

    286        @return true if the tree was fixed

    287     */

    288     private boolean bubbleUpFix(Node child)

    289     {

    290        if (child.color == NEGATIVE_RED) { fixNegativeRed(child); return true; }

    291        else if (child.color == RED)

    292        {

    293           if (child.left != null && child.left.color == RED)

    294           {

    295              fixDoubleRed(child.left); return true;

    296           }

    297           if (child.right != null && child.right.color == RED)

    298           {

    299              fixDoubleRed(child.right); return true;

    300           }

    301        }

    302        return false;

    303     }

    304

    305     /**

    306        Fixes a "double red" violation.

    307        @param child the child with a red parent

    308     */

    309     private void fixDoubleRed(Node child)

    310     {

    311        Node parent = child.parent;

    312        Node grandParent = parent.parent;

    313        if (grandParent == null) { parent.color = BLACK; return; }

    314        Node n1, n2, n3, t1, t2, t3, t4;

    315        if (parent == grandParent.left)

    316        {

    317           n3 = grandParent; t4 = grandParent.right;

    318           if (child == parent.left)

    319           {

    320              n1 = child; n2 = parent;

    321              t1 = child.left; t2 = child.right; t3 = parent.right;

    322           }

    323           else

    324           {

    325              n1 = parent; n2 = child;

    326              t1 = parent.left; t2 = child.left; t3 = child.right;

    327           }

    328        }

    329        else

    330        {

    331           n1 = grandParent; t1 = grandParent.left;

    332           if (child == parent.left)

    333           {

    334              n2 = child; n3 = parent;

    335              t2 = child.left; t3 = child.right; t4 = parent.right;

    336           }

    337           else

    338           {

    339              n2 = parent; n3 = child;

    340              t2 = parent.left; t3 = child.left; t4 = child.right;

    341           }

    342        }

    343

    344        replaceWith(grandParent, n2);

    345        n1.setLeftChild(t1);

    346        n1.setRightChild(t2);

    347        n2.setLeftChild(n1);

    348        n2.setRightChild(n3);

    349        n3.setLeftChild(t3);

    350        n3.setRightChild(t4);

    351        n2.color = grandParent.color - 1;

    352        n1.color = BLACK;

    353        n3.color = BLACK;

    354

    355        if (n2 == root)

    356        {

    357           root.color = BLACK;

    358        }

    359        else if (n2.color == RED && n2.parent.color == RED)

    360        {

    361           fixDoubleRed(n2);

    362        }

    363     }

    364

    365     /**

    366        Fixes a "negative red" violation.

    367        @param negRed the negative red node

    368     */

    369     private void fixNegativeRed(Node negRed)

    370     {

    371        Node parent = negRed.parent;

    372        Node child;

    373        if (parent.left == negRed)

    374        {

    375           Node n1 = negRed.left;

    376           Node n2 = negRed;

    377           Node n3 = negRed.right;

    378           Node n4 = parent;

    379           Node t1 = n3.left;

    380           Node t2 = n3.right;

    381           Node t3 = n4.right;

    382           n1.color = RED;

    383           n2.color = BLACK;

    384           n4.color = BLACK;

    385

    386           replaceWith(n4, n3);

    387           n3.setLeftChild(n2);

    388           n3.setRightChild(n4);

    389           n2.setLeftChild(n1);

    390           n2.setRightChild(t1);

    391           n4.setLeftChild(t2);

    392           n4.setRightChild(t3);

    393

    394           child = n1;

    395        }

    396        else // Mirror image

    397        {

    398           Node n4 = negRed.right;

    399           Node n3 = negRed;

    400           Node n2 = negRed.left;

    401           Node n1 = parent;

    402           Node t3 = n2.right;

    403           Node t2 = n2.left;

    404           Node t1 = n1.left;

    405           n4.color = RED;

    406           n3.color = BLACK;

    407           n1.color = BLACK;

    408

    409           replaceWith(n1, n2);

    410           n2.setRightChild(n3);

    411           n2.setLeftChild(n1);

    412           n3.setRightChild(n4);

    413           n3.setLeftChild(t3);

    414           n1.setRightChild(t2);

    415           n1.setLeftChild(t1);

    416

    417           child = n4;

    418        }

    419

    420        if (child.left != null && child.left.color == RED)

    421        {

    422           fixDoubleRed(child.left);

    423        }

    424        else if (child.right != null && child.right.color == RED)

    425        {

    426           fixDoubleRed(child.right);

    427        }

    428     }

    429  }

    430

    431

    432